__author__ = 'zoulida'
import  backtrader as bt
import pandas as pd
import numpy as np
import tushare as ts

def get_data(code,start='2010-01-01',end='2020-08-31'):
    df=ts.get_k_data(code,autype='qfq',start=start,end=end)
    df.index=pd.to_datetime(df.date)
    df['openinterest']=0
    df=df[['open','high','low','close','volume','openinterest']]
    return df

class TrendBand(bt.Indicator):
    lines = ('mid','top','bot',)
    params = (('maperiod',20),
              ('period',3),
              ('highRate',1.2),
              ('lowRate',0.85),)
    #与价格在同一张图
    plotinfo = dict(subplot=False)
    def __init__(self):
        ema = bt.ind.EMA(self.data, period=self.p.maperiod)
        #计算上中下轨线
        self.l.mid=bt.ind.EMA(ema,period=self.p.period)
        self.l.top=bt.ind.EMA(self.mid*self.p.highRate, \
                              period=self.p.period)
        self.l.bot=bt.ind.EMA(self.mid*self.p.lowRate, \
                              period=self.p.period)
        super(TrendBand, self).__init__()



class TestStrategy2(bt.Strategy):
    def __init__(self):
        TrendBand(self.data)

def main(data,strategy,pf=False):
    cerebro = bt.Cerebro()
    feed = bt.feeds.PandasData(dataname=data)
    cerebro.adddata(feed)
    #加载策略
    cerebro.addstrategy(strategy)
    # 设置初始资本为10,000
    startcash = 100000
    cerebro.broker.setcash(startcash)
    # 设置交易手续费为 0.1%
    cerebro.broker.setcommission(commission=0.001)
    cerebro.run()
    #获取回测结束后的总资金
    portvalue = cerebro.broker.getvalue()
    pnl = portvalue - startcash
    if pf:
        print(f'期初资金: 100000')
        print(f'总资金: {round(portvalue,2)}')
        print(f'净收益: {round(pnl,2)}')
    cerebro.plot()

data=get_data('601318','2010-01-01')
main(data,TestStrategy2)


class MyStrategy(bt.Strategy):
    params=(('period',20),)
    def __init__(self):
        self.order = None
        self.mid = TrendBand(self.data).mid
        self.top = TrendBand(self.data).top
        self.bot = TrendBand(self.data).bot
        #设置买入信号
        self.buy_sig=bt.And( \
            self.data.close>self.mid, \
            self.data.volume==bt.ind.Highest( \
                self.data.volume,period=self.p.period))
        #卖出信号
        self.sell_sig=self.data.close>self.top
    def next(self):
        if not self.position:
            # 得到当前的账户价值
            total_value = self.broker.getvalue()
            #1手=100股，满仓买入
            ss=int((total_value/100)/self.datas[0].close[0])*100
            if self.buy_sig:
                self.order=self.buy(size=ss)
        else:
            if self.sell_sig:
                self.close()


data=get_data('601318','2010-01-01')
main(data,MyStrategy,True)

